#include "hb_msgs/msg/heart_beat.hpp"
#include "rclcpp/rclcpp.hpp"
#include "sensor_msgs/image_encodings.hpp"
#include "sensor_msgs/msg/image.hpp"
#include "vaapi_decoder.hpp"

class decoder_node : public rclcpp::Node
{
  private:
	std::string input_url;
	std::string render_path;
	std::string scale_size;
	int width, height, size;
	int rotation;

	std::shared_ptr<vaapi_decoder> dec;

	rclcpp::Publisher<sensor_msgs::msg::Image>::SharedPtr img_pub;
	std::thread decode_thread;

	rclcpp::Publisher<hb_msgs::msg::HeartBeat>::SharedPtr hb_pub;
	rclcpp::TimerBase::SharedPtr hb_timer;

	int error_code = 0;
	int frame_cnt = 0;

	void decode_thread_func(void)
	{
		int ret;
		sensor_msgs::msg::Image img;
		sensor_msgs::msg::Image temp;
		img.encoding = sensor_msgs::image_encodings::TYPE_8UC3;
		img.is_bigendian = true;
		img.header.frame_id = "camera";
		RCLCPP_INFO(this->get_logger(), "decode thread running");
		while (rclcpp::ok())
		{
			img.data.resize(this->size);
			temp.data.resize(this->size);
			ret = this->dec->decode_frame(temp.data.data(), size, (int *)&temp.width,
										  (int *)&temp.height, (int *)&temp.step);
			if (ret < 0)
			{
				this->error_code = ret;
				RCLCPP_WARN(this->get_logger(),
							"error during decode, now try to open the stream again");
				this->dec->close_stream();
				this->dec->open_stream(this->input_url, this->render_path, this->scale_size);
				RCLCPP_WARN(this->get_logger(), "stream open, try resume streaming");
				continue;
			}

			int i, j;
			switch (this->rotation)
			{

				case 1: // anticlockwise 90
					img.width = temp.height;
					img.height = temp.width;
					img.step = img.width * 3;
					for (i = 0; i < img.height; i++)
					{
						for (j = 0; j < img.width; j++)
						{
							memcpy(img.data.data() + (i*img.width+j) * 3,
								   temp.data.data() + (j * temp.width + temp.width-1-i) * 3, 3);
						}
					}
					break;
				case 2: // rotate 180
					img.width = temp.width;
					img.height = temp.height;
					img.step = temp.step;
					for (i = 0; i < img.height; i++)
					{
						for (j = 0; j < img.width; j++)
						{
							memcpy(img.data.data() + (i * img.width + j) * 3,
								   temp.data.data() + ((temp.height - 1 - i) * temp.width +
													   temp.width - 1 - j) *
														  3,
								   3);
						}
					}
					break;
				case 3: // anticlockwise rotate 270, or clockwise 90
					img.width = temp.height;
					img.height = temp.width;
					img.step = img.width * 3;
					for (i = 0; i < img.height; i++)
					{
						for (j = 0; j < img.width; j++)
						{
							memcpy(img.data.data() + (i * img.width + j) * 3,
								   temp.data.data() + ((temp.height - 1 - j) * temp.width + i) * 3, 3);
						}
					}
					break;
				default: // no rotation
					img.width = temp.width;
					img.height = temp.height;
					img.step = temp.step;
					img.data = temp.data;
					break;
			}

			this->error_code = 0;
			this->frame_cnt += 1;
			RCLCPP_DEBUG(this->get_logger(), "frame decoded: %dx%d, step: %d", img.width,
						 img.height, img.step);
			img.header.stamp = rclcpp::Clock().now();
			this->img_pub->publish(img);
			RCLCPP_DEBUG(this->get_logger(), "published image");
		}
		RCLCPP_INFO(this->get_logger(), "decode thread exit, rclcpp not ok");
		this->error_code = 1;
	}

	void hb_callback(void)
	{
		hb_msgs::msg::HeartBeat hb;
		hb.name = this->get_fully_qualified_name();
		hb.stamp = rclcpp::Clock().now();
		switch (this->error_code)
		{
			case 0:
				hb.code = 0;
				hb.reason = "ok";
				break;
			case -1:
				hb.code = -1;
				hb.reason = "error with packets";
				break;
			case 1:
				hb.code = -3;
				hb.reason = "rclcpp shutdown";
				break;
			default:
				hb.code = -2;
				hb.reason = "error with frames";
		}
		if (hb.code == 0 && this->frame_cnt < 4)
		{
			hb.code = -4;
			hb.reason = "critical frame loss";
		}
		this->frame_cnt = 0;
		this->hb_pub->publish(hb);
	}

  public:
	decoder_node(std::string name) : Node(name)
	{
		RCLCPP_INFO(this->get_logger(), "node started");
		this->declare_parameter("input_url", "/home/anoki/ffmpeg/workspace/cat.mp4");
		this->declare_parameter("render_path", "/dev/dri/renderD128");
		this->declare_parameter("scale_size", "320x320");
		this->declare_parameter("rotation", 0);
		this->get_parameter("input_url", this->input_url);
		this->get_parameter("render_path", this->render_path);
		this->get_parameter("scale_size", this->scale_size);
		this->get_parameter("rotation", this->rotation);

		RCLCPP_INFO(this->get_logger(), "input_url: %s, render_path: %s, scale_size: %s, rotation: %d",
					this->input_url.c_str(), this->render_path.c_str(),
					this->scale_size.c_str(), this->rotation);
		sscanf(this->scale_size.c_str(), "%dx%d", &this->width, &this->height);
		this->size = this->width * this->height * 3;

		this->dec = std::make_shared<vaapi_decoder>();
		if (this->dec->open_stream(this->input_url, this->render_path, this->scale_size) < 0)
		{
			RCLCPP_ERROR(this->get_logger(),
						 "failed to open stream with given hardware device");
			return;
		}

		this->img_pub = this->create_publisher<sensor_msgs::msg::Image>("decoded_image", 1);

		this->decode_thread = std::thread(std::bind(&decoder_node::decode_thread_func, this));
		this->decode_thread.detach();

		RCLCPP_INFO(this->get_logger(),
					"decoder initilized, publishing image to topic \"decoded_image\"");

		this->hb_pub = this->create_publisher<hb_msgs::msg::HeartBeat>("/heartbeat", 1);
		this->hb_timer = this->create_wall_timer(std::chrono::milliseconds(200),
												 std::bind(&decoder_node::hb_callback, this));
	}

	~decoder_node()
	{
		this->dec->close_stream();
	}
};

int main(int argc, char **argv)
{
	rclcpp::init(argc, argv);
	auto node = std::make_shared<decoder_node>("decoder_node");
	rclcpp::spin(node);
	rclcpp::shutdown();
	return 0;
}
